package com.zdb.algorithm.kruskal;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class KruskalDemo {
    public static void main(String[] args) {

        Graph graph = createGraph();

        graph.print();

        System.out.println(graph.getEdgeList());
        // 从小到大排序
        graph.sortEdge();
        System.out.println(graph.getEdgeList());

        List<Graph.Edge> edgeList = graph.getEdgeList();
        int[] ends = new int[graph.getVertexSize()];

        // 最小连接树
        List<Graph.Edge> minEdges = new ArrayList<>();
        for(Graph.Edge edge : edgeList) {
            // 判断当前两个顶点的终点是否一致
            int eStart = graph.getEnd(ends, graph.getVertexIndex(edge.getStart()));
            int eEnd = graph.getEnd(ends, graph.getVertexIndex(edge.getEnd()));
            if(eStart != eEnd) {
                // 两个顶点的终点不一样，添加该边
               minEdges.add(edge);
                int eIdx = graph.getVertexIndex(edge.getEnd());
                int sIdx = graph.getVertexIndex(edge.getStart());
                int tmpIdx = 0;
                if(edge.getStart() > edge.getEnd()) {
                    // 交换开始结束
                    tmpIdx = eIdx;
                    eIdx = sIdx;
                    sIdx = tmpIdx;

                }
                // 更新最新的终点
                // 该点还没有边
                if(ends[eIdx] == 0) {
                    ends[eIdx] = eIdx;
                }
                ends[sIdx] = ends[eIdx];
            } else {
                // 两个顶点的终点一样，放弃该条边
            }
        }

        System.out.println("最小连接树：" + minEdges.size() + "," + minEdges);

    }

    public static Graph createGraph() {
        char[] vertexs = new char[]{'A', 'B', 'C', 'D', 'E', 'F', 'G'};
        Graph graph = new Graph(vertexs);

        // 增加边
        graph.addEdge('A', 'B', 12);
        graph.addEdge('A', 'F', 16);
        graph.addEdge('A', 'G', 14);
        graph.addEdge('B', 'C', 10);
        graph.addEdge('B', 'F', 7);
        graph.addEdge('C', 'D', 3);
        graph.addEdge('C', 'F', 6);
        graph.addEdge('C', 'E', 5);
        graph.addEdge('D', 'E', 4);
        graph.addEdge('E', 'F', 2);
        graph.addEdge('E', 'G', 8);
        graph.addEdge('F', 'G', 9);

        return graph;
    }
}

class Graph {
    /**
     顶点
     */
    char[] vertexs;
    /**
     * 存放边及权重
      */
    int[][] weight;
    /**
     顶点与vertexs中下标的对应关系
     */
    Map<Character, Integer> vertex2index;

    private List<Edge> edgeList;

    public static final int INF = Integer.MAX_VALUE;

    public Graph(char[] vertexs) {
        this.vertexs = vertexs;
        vertex2index = new HashMap<>();
        for (int i = 0; i < vertexs.length; i++) {
            vertex2index.put(vertexs[i], i);
        }
        weight = new int[vertexs.length][vertexs.length];
        for(int[] link : weight) {
            for (int i = 0; i < link.length; i++) {
                link[i] = INF;
            }
        }
        edgeList = new ArrayList<>();
    }

    public int getVertexSize() {
        return vertexs.length;
    }

    /**
     * 增加边
     * @param a
     * @param b
     * @param weight
     */
    public void addEdge(char a, char b, int weight) {
        int aIdx = vertex2index.get(a);
        int bIdx = vertex2index.get(b);
        if(this.weight[aIdx][bIdx] == INF && weight != INF) {
            this.weight[aIdx][bIdx] = weight;
            this.weight[bIdx][aIdx] = weight;
            edgeList.add(new Edge(a, b, weight));
        }
    }

    /**
     * 将边排序，根据weight 从小到大排序
     */
    public void sortEdge() {
        edgeList.sort(new Comparator<Edge>() {
            @Override public int compare(Edge o1, Edge o2) {
                return o1.weight - o2.weight;
            }
        });
    }

    /**
     * 获取所有的边
     * @return
     */
    public List<Edge> getEdgeList() {
        return edgeList;
    }

    /**
     * 获取终点
     * @param ends 长度为顶点个数，下标为顶点的下标，值为下标对应的顶点的终点
     * @param i 顶点对应的下标
     * @return
     */
    public int getEnd(int[] ends, int i) {
        // 循环找i位置的终点
        // ends[i] != 0 说明这个点已经有边了，
        // ends[i] != i 说明这个点的终点不是自己
        // 那么就去找ends[i]所指向的点(终点)，直到最后一个
        while(ends[i] != 0 && ends[i] != i) {
            i = ends[i];
        }
        return i;
    }

    public int getEdgeNum() {
        return edgeList.size();
    }

    /**
     * 获取顶点的下标
     * @param v
     * @return
     */
    public int getVertexIndex(char v) {
        return vertex2index.get(v);
    }

    /**
     * 通过下标获取顶点
     * @param index
     * @return
     */
    public char getVertexByIndex(int index) {
        return vertexs[index];
    }

    public void print() {
        for(int[] link : weight) {
            System.out.print("[");
            for (int i = 0; i < link.length; i++) {
                if(link[i]==INF) {
                    System.out.printf(" %5s", "INF");
                } else {
                    System.out.printf(" %5d", link[i]);
                }
                if(i != link.length -1) {
                    System.out.print(",");
                }
            }
            System.out.println("]");
        }
    }

    static class Edge {
        private char start;
        private char end;
        private int weight;

        public Edge(char start, char end, int weight) {
            this.start = start;
            this.end = end;
            this.weight = weight;
        }

        public char getStart() {
            return start;
        }

        public char getEnd() {
            return end;
        }

        public int getWeight() {
            return weight;
        }

        @Override public String toString() {
            return "[" + start +
                    "-" + end +
                    "(" + weight +
                    ")]";
        }
    }
}


